import asyncio
import base64
import logging
from pathlib import Path
from typing import Optional, Callable
import os

from .models import ProcessResult

logger = logging.getLogger(__name__)

class FaceSwapHandler:
    """
    换脸处理器 - 简化版
    固定使用extreme_beauty模式（blend=100）
    """
    
    def __init__(self):
        self.python_path = "/root/FaceFusion/facefusion_env/bin/python3"
        self.facefusion_path = "/root/FaceFusion/facefusion.py"
        self.temp_base_dir = Path("/root/FaceFusion/temp")
        self.temp_base_dir.mkdir(parents=True, exist_ok=True)
        self.timeout = 180
        
        logger.info("FaceSwapHandler initialized (extreme_beauty mode only)")
    
    async def process_faceswap(
        self,
        task_id: str,
        source_image_base64: str,
        target_image_base64: str,
        output_quality: int = 100,
        progress_callback: Optional[Callable[[int], None]] = None
    ) -> ProcessResult:
        """
        执行换脸处理（固定extreme_beauty模式）
        
        Args:
            task_id: 任务ID
            source_image_base64: 源图片base64编码
            target_image_base64: 目标图片base64编码
            output_quality: 输出图片质量 (1-100)
            progress_callback: 进度回调函数
            
        Returns:
            ProcessResult: 包含结果图片路径或错误信息的对象
        """
        try:
            logger.info(f"开始处理换脸任务: {task_id}, 模式: extreme_beauty (blend=100)")
            
            # 创建任务临时目录
            task_dir = self.temp_base_dir / task_id
            task_dir.mkdir(parents=True, exist_ok=True)
            
            # 解码并保存源图片和目标图片
            source_path = task_dir / "source.jpg"
            target_path = task_dir / "target.jpg"
            output_path = task_dir / "output.jpg"
            
            # 解码base64图片
            source_data = self._decode_base64_image(source_image_base64)
            target_data = self._decode_base64_image(target_image_base64)
            
            # 保存图片
            source_path.write_bytes(source_data)
            target_path.write_bytes(target_data)
            
            logger.info(f"图片已保存: source={source_path}, target={target_path}")
            
            # 构建命令（固定extreme_beauty配置）
            cmd = self._build_facefusion_command(
                source_path, target_path, output_path, output_quality
            )
            
            logger.info(f"执行命令: {' '.join(cmd)}")
            
            # 执行换脸处理
            if progress_callback:
                progress_callback(10)
            
            await self._execute_facefusion(cmd, progress_callback)
            
            # 验证输出文件
            if not output_path.exists():
                raise RuntimeError(f"输出文件未生成: {output_path}")
            
            logger.info(f"换脸任务完成: {task_id}, 输出: {output_path}")
            
            return ProcessResult(success=True, output_path=str(output_path))
            
        except Exception as e:
            logger.error(f"换脸处理失败 - 任务ID: {task_id}, 错误: {str(e)}")
            return ProcessResult(success=False, error_message=str(e))
    
    def _decode_base64_image(self, base64_str: str) -> bytes:
        """解码base64图片"""
        # 移除data:image前缀
        if ',' in base64_str:
            base64_str = base64_str.split(',')[1]
        return base64.b64decode(base64_str)
    
    def _build_facefusion_command(
        self,
        source_path: Path,
        target_path: Path,
        output_path: Path,
        quality: int
    ) -> list:
        """
        构建facefusion命令（固定extreme_beauty配置）
        对应脚本配置: blend=100, inswapper_128, gfpgan_1.4
        """
        cmd = [
            self.python_path,
            self.facefusion_path,
            "headless-run",
            "-s", str(source_path),
            "-t", str(target_path),
            "-o", str(output_path),
            "--processors", "face_swapper", "face_enhancer",
            "--face-swapper-model", "inswapper_128",
            "--face-enhancer-model", "gfpgan_1.4",
            "--face-enhancer-blend", "100",  # extreme_beauty: 最强美颜
            "--face-detector-model", "yolo_face",
            "--face-landmarker-model", "2dfan4",
            "--face-selector-mode", "one",
            "--output-image-quality", str(quality)
        ]
        
        return cmd
    
    async def _execute_facefusion(
        self,
        cmd: list,
        progress_callback: Optional[Callable[[int], None]] = None
    ) -> None:
        """异步执行facefusion命令"""
        try:
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=self.temp_base_dir.parent
            )
            
            # 模拟进度更新
            if progress_callback:
                for progress in [30, 50, 70, 90]:
                    await asyncio.sleep(2)
                    progress_callback(progress)
            
            stdout, stderr = await asyncio.wait_for(
                process.communicate(),
                timeout=self.timeout
            )
            
            if process.returncode != 0:
                error_msg = stderr.decode('utf-8', errors='ignore')
                raise RuntimeError(f"FaceFusion执行失败 (code {process.returncode}): {error_msg}")
            
            logger.info("FaceFusion命令执行成功")
            
        except asyncio.TimeoutError:
            process.kill()
            await process.wait()
            raise TimeoutError(f"换脸处理超时 ({self.timeout}秒)")
        except Exception as e:
            logger.error(f"执行FaceFusion命令时发生错误: {str(e)}")
            raise RuntimeError(f"FaceFusion执行失败: {str(e)}")

# 创建全局实例
faceswap_handler = FaceSwapHandler()
